import tensorflow as tf
import numpy as np
np.set_printoptions(formatter={'float': '{: 0.4f}'.format})
## 1. data prepare
minst = tf.keras.datasets.mnist
img_rows,img_cols = 28,28
(x_train, y_train), (x_test, y_test) = minst.load_data()
x_train = x_train.reshape(x_train.shape[0],img_rows,img_cols,1)
x_test = x_test.reshape(x_test.shape[0],img_rows,img_cols,1)

x_train = x_train.astype('float32')
x_test = x_test.astype('float32')
x_train = x_train / 255
x_test = x_test / 255
y_train_onehot = tf.keras.utils.to_categorical(y_train)
y_test_onehot = tf.keras.utils.to_categorical(y_test)

## 2. full precision
model = tf.keras.Sequential()
model = tf.keras.models.load_model('models/mnist_tf2_fw.h5')
score = model.evaluate(x_test, y_test_onehot, verbose=0)
print('full accuracy:', "{:.5f}".format(score[1]))

## 3. Lloyds函数
def Lloyds(level, quan):
    temp_sort_Q2 = np.zeros(temp_sort.size)
    quan2 = quan.copy()
    loss_temp = 0
    while True:
        edge = (quan2[0:-1] + quan2[1:]) / 2
        quan2[0] = temp_sort[(temp_sort <= edge[0])].mean()
        for kk in range(1, level - 1):
            quan2[kk] = temp_sort[(edge[kk - 1] < temp_sort) * (temp_sort <= edge[kk])].mean()
        quan2[level - 1] = temp_sort[(edge[level - 2] < temp_sort)].mean()
        # print("quan2 : ",quan2)
        for j in range(temp_sort.size):
            temp_sort_Q2[j] = quan2[np.argmin(abs(temp_sort[j] - quan2))]
        # print("Q2",temp_sort_Q2)
        # print("temp ",temp_sort)
        loss2 = np.square(temp_sort - temp_sort_Q2)
        # loss2 = np.abs(temp_sort-temp_sort_Q2)
        loss2 = loss2.mean()
        # print("loss2 : %f"%(loss2))
        if abs(loss2 - loss_temp) < 1e-5:
            break
        loss_temp = loss2
    return quan2

# 设置待量化的level点数，不是比特数
level= 4
print("level:",level)
## 4.1 逆采样量化
model = tf.keras.models.load_model('models/mnist_tf2_fw.h5')
for i in [0,2,5]:
    temp = model.layers[i].get_weights()[0]
    temp2=temp.reshape(1,-1)
    temp_sort=np.sort(temp2,axis=None)
    num_list = (np.round(np.arange(0.5/level,1,1/level)*temp.size)).astype(np.int32)
    quan = temp_sort[list(num_list)]
    # print(quan)
    #print(temp.size)
    for j in range(temp.size):
        temp2[0][j]=quan[np.argmin(abs(temp2[0][j]-quan))]
    temp_quan = temp2.reshape(temp.shape)
    model.layers[i].set_weights([temp_quan])
score = model.evaluate(x_test, y_test_onehot, verbose=0)
print('inv_Sampling:', "{:.5f}".format(score[1]))
## 4.2 Lloyd 量化
model = tf.keras.models.load_model('models/mnist_tf2_fw.h5')
for i in [0,2,5]:
    temp = model.layers[i].get_weights()[0]
    temp2=temp.reshape(1,-1)
    temp_sort=np.sort(temp2,axis=None)
    num_list = (np.round(np.arange(0.5/level,1,1/level)*temp.size)).astype(np.int32)
    quan = temp_sort[list(num_list)]
    #print(quan)
    quan2 = Lloyds(level,quan)
    # print(quan2)
    #print(temp.size)
    for j in range(temp.size):
        temp2[0][j]=quan2[np.argmin(abs(temp2[0][j]-quan2))]
    temp_quan = temp2.reshape(temp.shape)
    model.layers[i].set_weights([temp_quan])
score = model.evaluate(x_test, y_test_onehot, verbose=0)
print('Lloyd:', "{:.5f}".format(score[1]))

## 4.3 均匀量化
model = tf.keras.models.load_model('models/mnist_tf2_fw.h5')
for i in [0,2,5]:
    temp = model.layers[i].get_weights()[0]
    temp2=temp.reshape(1,-1)
    delta = (temp2.max()-temp2.min())/level
    quan = np.arange(temp2.min()+0.5*delta,temp2.max(),delta)
    #print(temp2.min(),temp2.max(),quan)
    for j in range(temp.size):
        temp2[0][j]=quan[np.argmin(abs(temp2[0][j]-quan))]
    temp_quan = temp2.reshape(temp.shape)
    model.layers[i].set_weights([temp_quan])
score = model.evaluate(x_test, y_test_onehot, verbose=0)
print('normal:', "{:.5f}".format(score[1]))